import os

import numpy as np
from osgeo import gdal
from osgeo.gdalconst import GDT_Float32


def write_tif(file_path, data, geotrans, projection, nodata, gdal_type):
    driver = gdal.GetDriverByName("GTiff")
    rows, cols = data.shape
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    dataset.SetGeoTransform(geotrans)
    # 定义投影
    # prj = osr.SpatialReference()
    # prj.ImportFromEPSG(4326)
    # dataset.SetProjection(prj.ExportToWkt())
    dataset.SetProjection(projection)
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    dataset.FlushCache()
    print("写入成功：{}".format(file_path))


def computed_mean_lst(day_lst_dir, night_lst_dir, output_dir):
    day_lst_list = [os.path.join(day_lst_dir, file) for file in os.listdir(day_lst_dir)]
    night_lst_list = [os.path.join(night_lst_dir, file) for file in os.listdir(night_lst_dir)]
    for day_lst_path, night_lst_path in zip(day_lst_list, night_lst_list):
        day_lst_dataset = gdal.Open(day_lst_path)
        night_lst_dataset = gdal.Open(night_lst_path)
        day_lst = day_lst_dataset.ReadAsArray() * 0.02
        night_lst = night_lst_dataset.ReadAsArray() * 0.02
        lst_geo = day_lst_dataset.GetGeoTransform()
        lst_pro = day_lst_dataset.GetProjection()
        day_lst[day_lst <= 0] = np.nan
        night_lst[night_lst <= 0] = np.nan
        mean_lst = (day_lst + night_lst) / 2
        mean_lst[mean_lst <= 0] = -9999
        mean_lst_path = os.path.join(output_dir, os.path.basename(day_lst_path).split("_")[2] + ".tif")
        write_tif(mean_lst_path, mean_lst, lst_geo, lst_pro, -9999, GDT_Float32)


if __name__ == '__main__':
    day_lst_dir = r"G:\test\SMAP_NEW\2015_day"
    night_lst_dir = r"G:\test\SMAP_NEW\2015_night"
    output_dir = r"G:\test\SMAP_NEW\2015_lst"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    computed_mean_lst(day_lst_dir, night_lst_dir, output_dir)
